import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image
import os

# 创建自定义的draw_on函数来修复np.int问题
def fixed_draw_on(img, faces):
    """
    修复版本的draw_on函数，将np.int替换为int
    """
    dimg = img.copy()
    for i in range(len(faces)):
        face = faces[i]
        # 修复：使用int而不是np.int
        box = face.bbox.astype(int)  
        color = (0, 0, 255)
        cv2.rectangle(dimg, (box[0], box[1]), (box[2], box[3]), color, 2)
        if face.kps is not None:
            kps = face.kps
            for l in range(kps.shape[0]):
                color = (0, 255, 0)
                cv2.circle(dimg, (int(kps[l, 0]), int(kps[l, 1])), 1, color, 2)
    return dimg

# 保存图片函数
def save_image(img, output_path, filename="output.jpg"):
    """
    保存图片到指定路径
    
    参数:
        img: 图像数据
        output_path: 输出路径
        filename: 文件名
    
    返回:
        保存的完整路径
    """
    # 确保输出路径存在
    os.makedirs(output_path, exist_ok=True)
    
    # 构建完整的输出路径
    full_path = os.path.join(output_path, filename)
    
    # 保存图片
    cv2.imwrite(full_path, img)
    
    return full_path

# 初始化人脸分析器
def init_face_analyzer(use_gpu=False):
    """
    初始化人脸分析器
    
    参数:
        use_gpu: 是否使用GPU
    
    返回:
        初始化好的FaceAnalysis对象
    """
    providers = ['CPUExecutionProvider']
    if use_gpu:
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    
    app = FaceAnalysis(providers=providers)
    ctx_id = 0 if use_gpu else -1
    app.prepare(ctx_id=ctx_id, det_size=(640, 640))
    
    return app

# 检测图片中的人脸
def detect_faces(image_path, use_gpu=False):
    """
    检测指定路径图片中的人脸
    
    参数:
        image_path: 图片路径
        use_gpu: 是否使用GPU
    
    返回:
        原始图像, 检测到的人脸列表
    """
    # 初始化人脸分析器
    app = init_face_analyzer(use_gpu)
    
    # 读取图片
    if os.path.exists(image_path):
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图片: {image_path}")
    else:
        # 如果是insightface内置的示例图片
        try:
            img = ins_get_image(image_path)
        except:
            raise ValueError(f"图片不存在: {image_path}")
    
    # 检测人脸
    faces = app.get(img)
    
    return img, faces

# 处理图片并保存结果
def image_detect_tag(image_path, output_dir="./output", output_filename=None, use_gpu=False):
    """
    处理图片并保存结果
    
    参数:
        image_path: 图片路径
        output_dir: 输出目录
        output_filename: 输出文件名，如果为None则自动生成
        use_gpu: 是否使用GPU
    
    返回:
        检测到的人脸数量, 保存的图片路径
    调用：
    # 示例1: 处理内置示例图片
        faces_count1, saved_path1 = image_detect_tag("t1", "./output", "example_output.jpg")
        print(f"示例1: 成功检测到 {faces_count1} 个人脸，结果已保存到 {saved_path1}")
    """
    try:
        # 检测人脸
        img, faces = detect_faces(image_path, use_gpu)
        
        # 如果未指定输出文件名，则根据输入文件名生成
        if output_filename is None:
            base_name = os.path.basename(image_path)
            name, ext = os.path.splitext(base_name)
            output_filename = f"{name}_detected{ext}"
        
        # 绘制人脸框
        result_img = fixed_draw_on(img, faces)
        
        # 保存结果
        saved_path = save_image(result_img, output_dir, output_filename)
        
        return len(faces), saved_path
    
    except Exception as e:
        print(f"处理图片时发生错误: {e}")
        import traceback
        traceback.print_exc()
        return 0, None

def main():
    try:
        # 示例2: 处理自定义路径的图片 (如果存在)
        custom_image = "./input_images/普京-图片.jpg"  # 替换为您自己的图片路径
        if os.path.exists(custom_image):
            faces_count2, saved_path2 = image_detect_tag(custom_image)
            print(f"示例2: 成功检测到 {faces_count2} 个人脸，结果已保存到 {saved_path2}")
        
        return True
    except Exception as e:
        print(f"发生错误: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = main()
    print(f"测试{'成功' if success else '失败'}") 